// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use thiserror::Error;

use xgpu_common::{api_name::ApiFuncName, ipc::message::Argument};

use super::handler::*;

#[derive(Debug, Error)]
pub enum ServerErr {
    #[error("API run failed, api: {0}, error code: {1}")]
    ApiRunError(String, i32),

    #[error("Invalid argument type: {0}")]
    InvalidType(String),
}

pub trait ApiHandler: Send + Sync {
    fn handle_api(&self, args: &mut [Argument<'_>]) -> Result<Argument<'static>, ServerErr>;
}

macro_rules! dispatch_api {
    (
        $id:expr,
        { $($enum_variant:path => $handler:ident),* $(,)? }
    ) => {
        match $id as i32 {
            $(
                $enum_variant => Ok(&$handler as &'static dyn ApiHandler),
            )*
            _ => Err(ServerErr::ApiRunError(
                "Unknown API method ID".to_string(),
                $id as i32,
            )),
        }
    };
}

pub struct Api;

impl Api {
    #[inline]
    fn get_handler(id: u64) -> Result<&'static dyn ApiHandler, ServerErr> {
        dispatch_api!(id, {
            ApiFuncName::FUNC_CUDADEVICESYNCHRONIZE => CudaDeviceSynchronizeHandler,
            ApiFuncName::FUNC_CUDADEVICEGETSTREAMPRIORITYRANGE => CudaDeviceGetStreamPriorityRangeHandler,
            ApiFuncName::FUNC_CUDAGETLASTERROR => CudaGetLastErrorHandler,
            ApiFuncName::FUNC_CUDAPEEKATLASTERROR => CudaPeekAtLastErrorHandler,
            ApiFuncName::FUNC_CUDAGETDEVICECOUNT => CudaGetDeviceCountHandler,
            ApiFuncName::FUNC_CUDAGETDEVICEPROPERTIES_V2 => CudaGetDevicePropertiesV2Handler,
            ApiFuncName::FUNC_CUDADEVICEGETATTRIBUTE => CudaDeviceGetAttributeHandler,
            ApiFuncName::FUNC_CUDASETDEVICE => CudaSetDeviceHandler,
            ApiFuncName::FUNC_CUDAGETDEVICE => CudaGetDeviceHandler,
            ApiFuncName::FUNC_CUDASTREAMCREATEWITHPRIORITY => CudaStreamCreateWithPriorityHandler,
            ApiFuncName::FUNC_CUDATHREADEXCHANGESTREAMCAPTUREMODE => CudaThreadExchangeStreamCaptureModeHandler,
            ApiFuncName::FUNC_CUDASTREAMISCAPTURING => CudaStreamIsCapturingHandler,
            ApiFuncName::FUNC_CUDASTREAMWAITEVENT => CudaStreamWaitEventHandler,
            ApiFuncName::FUNC_CUDASTREAMSYNCHRONIZE => CudaStreamSynchronizeHandler,
            ApiFuncName::FUNC_NCCLGETVERSION => NcclGetVersionHandler,
            ApiFuncName::FUNC_CUDAEVENTCREATEWITHFLAGS => CudaEventCreateWithFlagsHandler,
            ApiFuncName::FUNC_CUDAEVENTRECORD => CudaEventRecordHandler,
            ApiFuncName::FUNC_CUDAEVENTRECORDWITHFLAGS => CudaEventRecordWithFlagsHandler,
            ApiFuncName::FUNC_CUDAEVENTQUERY => CudaEventQueryHandler,
            ApiFuncName::FUNC_CUDAEVENTSYNCHRONIZE => CudaEventSynchronizeHandler,
            ApiFuncName::FUNC_CUDAEVENTDESTROY => CudaEventDestroyHandler,
            ApiFuncName::FUNC_CUDAEVENTELAPSEDTIME => CudaEventElapsedTimeHandler,
            ApiFuncName::FUNC_CUDAMEMSET => CudaMemsetHandler,
            ApiFuncName::FUNC_CUDAPOINTERGETATTRIBUTES => CudaPointerGetAttributesHandler,
            ApiFuncName::FUNC_CUDEVICEGET => CuDeviceGetHandler,
            ApiFuncName::FUNC_CUDEVICEGETATTRIBUTE => CuDeviceGetAttributeHandler,
            ApiFuncName::FUNC_CUDEVICEPRIMARYCTXGETSTATE => CuDevicePrimaryCtxGetStateHandler,
            ApiFuncName::FUNC_CUCTXGETCURRENT => CuCtxGetCurrentHandler,
            ApiFuncName::FUNC_CUFUNCGETATTRIBUTE => CuFuncGetAttributeHandler,
            ApiFuncName::FUNC_CUDAMEMSETASYNC => CudaMemsetAsyncHandler,
            ApiFuncName::FUNC_CUDAFUNCSETATTRIBUTE => CudaFuncSetAttributeHandler,
            ApiFuncName::FUNC_CUDAFUNCGETATTRIBUTES => CudaFuncGetAttributesHandler,
            ApiFuncName::FUNC_CUBLASCREATE_V2 => CublasCreateV2handler,
            ApiFuncName::FUNC_CUBLASSETWORKSPACE_V2 => CublasSetWorkspaceV2handler,
            ApiFuncName::FUNC_CUBLASSETSTREAM_V2 => CublasSetStreamV2handler,
            ApiFuncName::FUNC_CUBLASGETSTREAM_V2 => CublasGetStreamV2handler,
            ApiFuncName::FUNC_CUBLASLTMATMULPREFERENCECREATE => CublasLtMatmulPreferenceCreateHandler,
            ApiFuncName::FUNC_CUBLASLTMATMULDESCCREATE => CublasLtMatmulDescCreateHandler,
            ApiFuncName::FUNC_CUBLASLTMATMULDESCSETATTRIBUTE => CublasLtMatmulDescSetAttributeHandler,
            ApiFuncName::FUNC_CUBLASLTMATRIXLAYOUTCREATE => CublasLtMatrixLayoutCreateHandler,
            ApiFuncName::FUNC_CUBLASLTMATMULPREFERENCESETATTRIBUTE => CublasLtMatmulPreferenceSetAttributeHandler,
            ApiFuncName::FUNC_CUBLASGETMATHMODE => CublasGetMathModeHandler,
            ApiFuncName::FUNC_CUBLASSETMATHMODE => CublasSetMathModeHandler,
            ApiFuncName::FUNC_CUBLASLTMATMULALGOGETHEURISTIC => CublasLtMatmulAlgoGetHeuristicHandler,
            ApiFuncName::FUNC_CUBLASLTMATMUL => CublasLtMatmulHandler,
            ApiFuncName::FUNC_CUBLASLTMATRIXLAYOUTDESTROY => CublasLtMatrixLayoutDestroyHandler,
            ApiFuncName::FUNC_CUBLASLTMATMULDESCDESTROY => CublasLtMatmulDescDestroyHandler,
            ApiFuncName::FUNC_CUBLASLTMATMULPREFERENCEDESTROY => CublasLtMatmulPreferenceDestroyHandler,
            ApiFuncName::FUNC_CUBLASSGEMM_V2 => CublasSgemmV2handler,
            ApiFuncName::FUNC_CUDAMEMCPYASYNC => CudaMemcpyAsyncHandler,
            ApiFuncName::FUNC_CUDAMALLOC => CudaMallocHandler,
            ApiFuncName::FUNC_CUDAFREE => CudaFreeHandler,
            ApiFuncName::FUNC_NVMLINIT_V2 => NvmlInitV2Handler,
            ApiFuncName::FUNC_NVMLSHUTDOWN => NvmlShutdownHandler,
            ApiFuncName::FUNC_NVMLDEVICEGETCOUNT_V2 => NvmlDeviceGetCountV2handler,
            ApiFuncName::FUNC_NCCLCOMMDESTROY => NcclCommDestroyHandler,
            ApiFuncName::FUNC_NCCLCOMMABORT => NcclCommAbortHandler,
            ApiFuncName::FUNC_NCCLCOMMFINALIZE => NcclCommFinalizeHandler,
            ApiFuncName::FUNC_NCCLCOMMSPLIT => NcclCommSplitHandler,
            ApiFuncName::FUNC_NCCLCOMMGETASYNCERROR => NcclCommGetAsyncErrorHandler,
            ApiFuncName::FUNC_NCCLBCAST => NcclBcastHandler,
            ApiFuncName::FUNC_NCCLALLREDUCE => NcclAllReduceHandler,
            ApiFuncName::FUNC_NCCLALLGATHER => NcclAllGatherHandler,
            ApiFuncName::FUNC_NCCLSEND => NcclSendHandler,
            ApiFuncName::FUNC_NCCLRECV => NcclRecvHandler,
            ApiFuncName::FUNC_NCCLGROUPSTART => NcclGroupStartHandler,
            ApiFuncName::FUNC_NCCLGROUPEND => NcclGroupEndHandler,
            ApiFuncName::FUNC_NCCLGETUNIQUEID => NcclGetUniqueIdHandler,
            ApiFuncName::FUNC_NCCLCOMMINITRANKCONFIG => NcclCommInitRankConfigHandler,
            ApiFuncName::FUNC_CUDALAUNCHKERNEL => CudaLaunchKernelHandler,
            ApiFuncName::FUNC_NCCLREDUCE => NcclReduceHandler,
            ApiFuncName::FUNC_CUBLASGEMMEX => CublasGemmExHandler,
            ApiFuncName::FUNC_CUBLASSGEMMSTRIDEDBATCHED => CublasSgemmStridedBatchedHandler,
            ApiFuncName::FUNC_CUBLASGEMMSTRIDEDBATCHEDEX => CublasGemmStridedBatchedExHandler,
            ApiFuncName::FUNC_CUBLASLTCREATE => CublasLtCreateHandler,
            ApiFuncName::FUNC_CUBLASLTDESTROY => CublasLtDestroyHandler,
            ApiFuncName::FUNC_CUDAOCCUPANCYMAXACTIVEBLOCKSPERMULTIPROCESSORWITHFLAGS => CudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlagsHandler,
            ApiFuncName::FUNC_NCCLCOMMINITRANKCONFIGRECOVERY => NcclCommInitRankConfigRecoveryHandler,
            ApiFuncName::FUNCLOADDYNLIBRARY => LoadLibraryHandler,
        })
    }

    pub fn invoke(id: u64, args: &mut [Argument<'_>]) -> Result<Argument<'static>, ServerErr> {
        Self::get_handler(id)?.handle_api(args)
    }
}
